import pyaudio
from websocket import create_connection, WebSocketConnectionClosedException
import websocket
import json
import time
import threading
import hashlib
import hmac
import base64
from urllib.parse import quote

class RealTimeASRClient:
    def __init__(self, base_url, app_id, api_key, lang="cn", punc="1", pd=None, vad_mdn=1, role_type=None, eng_lang_type=None, transType=None, transStrategy=None, targetLang=None):
        ts = str(int(time.time()))
        base_string = (app_id + ts).encode('utf-8')
        md5 = hashlib.md5()
        md5.update(base_string)
        base_string_md5 = md5.hexdigest().encode('utf-8')

        api_key_bytes = api_key.encode('utf-8')
        signa = hmac.new(api_key_bytes, base_string_md5, hashlib.sha1).digest()
        signa = base64.b64encode(signa).decode('utf-8')

        params = {
            "appid": app_id,
            "ts": ts,
            "signa": signa,
            "lang": lang,
            "punc": punc
        }
        if pd:
            params["pd"] = pd
        if vad_mdn != 1:
            params["vadMdn"] = vad_mdn
        if role_type:
            params["roleType"] = role_type
        if eng_lang_type:
            params["engLangType"] = eng_lang_type
        if transType:
            params["transType"] = transType
        if transStrategy:
            params["transStrategy"] = transStrategy
        if targetLang:
            params["targetLang"] = targetLang

        url = base_url + "?" + "&".join([f"{k}={quote(str(v))}" for k, v in params.items()])

        try:
            self.ws = create_connection(url)
        except Exception as e:
            print(f"WebSocket 连接失败: {e}")
            self.ws = None

        if self.ws:
            self.trecv = threading.Thread(target=self.recv)
            self.trecv.start()

    def send_realtime(self, audio_stream):
        if not self.ws:
            return

        try:
            while True:
                chunk = audio_stream.read(1280)
                if not chunk:
                    break
                self.ws.send(chunk)
                time.sleep(0.04)  # 控制发送频率
        finally:
            if self.ws:
                self.ws.send(json.dumps({"end": True}).encode('utf-8'))  # 发送结束标志

    def recv(self):
        full_text = []  # 存储完整的转写结果
        try:
            while self.ws.connected:
                result = self.ws.recv()
                if len(result) == 0:
                    break
                result_dict = json.loads(result)
                if result_dict["action"] == "result":
                    text = self.parse_result(result_dict)
                    if text.strip():
                        full_text.append(text)
        except WebSocketConnectionClosedException:
            pass

        if full_text:
            print("".join(full_text))  # 输出最终转写结果

    def parse_result(self, result_dict):
        try:
            data_str = result_dict.get("data", "{}")
            data_dict = json.loads(data_str)

            cn = data_dict.get("cn", {}).get("st", {}).get("rt", [])
            words = []
            for item in cn:
                ws = item.get("ws", [])
                for word_group in ws:
                    cw = word_group.get("cw", [])
                    for word_info in cw:
                        word = word_info.get("w", "").strip()
                        if word:
                            words.append(word)

            return "".join(words)
        except Exception:
            return ""

    def close(self):
        if self.ws:
            self.ws.close()

if __name__ == '__main__':
    APP_ID = "3c4e1474"  # 替换为你的 APPID
    API_KEY = "7ee110d694de56a0af762bdb9721d3de"  # 替换为你的 APIKey
    BASE_URL = "ws://rtasr.xfyun.cn/v1/ws"

    CHUNK = 1024
    FORMAT = pyaudio.paInt16
    CHANNELS = 1
    RATE = 16000
    DEVICE_INDEX = 1  # 替换为你的麦克风设备索引

    p = pyaudio.PyAudio()
    stream = p.open(format=FORMAT,
                    channels=CHANNELS,
                    rate=RATE,
                    input=True,
                    frames_per_buffer=CHUNK,
                    input_device_index=DEVICE_INDEX)

    client = RealTimeASRClient(BASE_URL, APP_ID, API_KEY, lang="cn", punc="1", pd="edu", vad_mdn=1, role_type=None, eng_lang_type=None)
    if client.ws:
        client.send_realtime(stream)